{
 "cells": [
  {
   "cell_type": "markdown",
   "source": "# Building AI Agents with AWS Bedrock and Koog Framework\n\nWelcome to this comprehensive guide on creating intelligent AI agents using the Koog framework with AWS Bedrock integration. In this notebook, we'll walk through building a functional agent that can control a simple switch device through natural language commands.\n\n## What You'll Learn\n\n- How to define custom tools for AI agents using Kotlin annotations\n- Setting up AWS Bedrock integration for LLM-powered agents\n- Creating tool registries and connecting them to agents\n- Building interactive agents that can understand and execute commands\n\n## Prerequisites\n\n- AWS Bedrock access with appropriate permissions\n- AWS credentials configured (access key and secret key)\n- Basic understanding of Kotlin coroutines\n\nLet's dive into building our first Bedrock-powered AI agent!",
   "metadata": {}
  },
  {
   "cell_type": "code",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-08-18T21:24:27.587338Z",
     "start_time": "2025-08-18T21:24:27.551913Z"
    }
   },
   "source": [
    "%useLatestDescriptors\n",
    "// %use koog"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:33:01.678131Z",
     "start_time": "2025-08-18T21:33:01.565343Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.agents.core.tools.annotations.LLMDescription\n",
    "import ai.koog.agents.core.tools.annotations.Tool\n",
    "import ai.koog.agents.core.tools.reflect.ToolSet\n",
    "\n",
    "// Simple state-holding device that our agent will control\n",
    "class Switch {\n",
    "    private var state: Boolean = false\n",
    "\n",
    "    fun switch(on: Boolean) {\n",
    "        state = on\n",
    "    }\n",
    "\n",
    "    fun isOn(): Boolean {\n",
    "        return state\n",
    "    }\n",
    "}\n",
    "\n",
    "/**\n",
    " * ToolSet implementation that exposes switch operations to the AI agent.\n",
    " *\n",
    " * Key concepts:\n",
    " * - @Tool annotation marks methods as callable by the agent\n",
    " * - @LLMDescription provides natural language descriptions for the LLM\n",
    " * - ToolSet interface allows grouping related tools together\n",
    " */\n",
    "class SwitchTools(val switch: Switch) : ToolSet {\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Switches the state of the switch to on or off\")\n",
    "    fun switchState(state: Boolean): String {\n",
    "        switch.switch(state)\n",
    "        return \"Switch turned ${if (state) \"on\" else \"off\"} successfully\"\n",
    "    }\n",
    "\n",
    "    @Tool\n",
    "    @LLMDescription(\"Returns the current state of the switch (on or off)\")\n",
    "    fun getCurrentState(): String {\n",
    "        return \"Switch is currently ${if (switch.isOn()) \"on\" else \"off\"}\"\n",
    "    }\n",
    "}"
   ],
   "outputs": [],
   "execution_count": 1
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:33:41.533575Z",
     "start_time": "2025-08-18T21:33:41.419083Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.agents.core.tools.ToolRegistry\n",
    "import ai.koog.agents.core.tools.reflect.asTools\n",
    "\n",
    "// Create our switch instance\n",
    "val switch = Switch()\n",
    "\n",
    "// Build the tool registry with our switch tools\n",
    "val toolRegistry = ToolRegistry {\n",
    "    // Convert our ToolSet to individual tools and register them\n",
    "    tools(SwitchTools(switch).asTools())\n",
    "}\n",
    "\n",
    "println(\"✅ Tool registry created with ${toolRegistry.tools.size} tools:\")\n",
    "toolRegistry.tools.forEach { tool ->\n",
    "    println(\"  - ${tool.name}\")\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Tool registry created with 2 tools:\n",
      "  - getCurrentState\n",
      "  - switchState\n"
     ]
    }
   ],
   "execution_count": 3
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:35:21.615537Z",
     "start_time": "2025-08-18T21:35:21.550364Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.prompt.executor.clients.bedrock.BedrockClientSettings\n",
    "import ai.koog.prompt.executor.clients.bedrock.BedrockRegions\n",
    "\n",
    "val region = BedrockRegions.US_WEST_2.regionCode\n",
    "val maxRetries = 3\n",
    "\n",
    "// Configure Bedrock client settings\n",
    "val bedrockSettings = BedrockClientSettings(\n",
    "    region = region, // Choose your preferred AWS region\n",
    "    maxRetries = maxRetries // Number of retry attempts for failed requests\n",
    ")\n",
    "\n",
    "println(\"🌐 Bedrock configured for region: $region\")\n",
    "println(\"🔄 Max retries set to: $maxRetries\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🌐 Bedrock configured for region: us-west-2\n",
      "🔄 Max retries set to: 3\n"
     ]
    }
   ],
   "execution_count": 4
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:35:58.807364Z",
     "start_time": "2025-08-18T21:35:58.601949Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.prompt.executor.llms.all.simpleBedrockExecutor\n",
    "\n",
    "// Create the Bedrock LLM executor with credentials from environment\n",
    "val executor = simpleBedrockExecutor(\n",
    "    awsAccessKeyId = System.getenv(\"AWS_BEDROCK_ACCESS_KEY\")\n",
    "        ?: throw IllegalStateException(\"AWS_BEDROCK_ACCESS_KEY environment variable not set\"),\n",
    "    awsSecretAccessKey = System.getenv(\"AWS_BEDROCK_SECRET_ACCESS_KEY\")\n",
    "        ?: throw IllegalStateException(\"AWS_BEDROCK_SECRET_ACCESS_KEY environment variable not set\"),\n",
    "    settings = bedrockSettings\n",
    ")\n",
    "\n",
    "println(\"🔐 Bedrock executor initialized successfully\")\n",
    "println(\"💡 Pro tip: Set AWS_BEDROCK_ACCESS_KEY and AWS_BEDROCK_SECRET_ACCESS_KEY environment variables\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔐 Bedrock executor initialized successfully\n",
      "💡 Pro tip: Set AWS_BEDROCK_ACCESS_KEY and AWS_BEDROCK_SECRET_ACCESS_KEY environment variables\n"
     ]
    }
   ],
   "execution_count": 5
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:36:11.490818Z",
     "start_time": "2025-08-18T21:36:11.395246Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import ai.koog.agents.core.agent.AIAgent\n",
    "import ai.koog.prompt.executor.clients.bedrock.BedrockModels\n",
    "\n",
    "val agent = AIAgent(\n",
    "    executor = executor,\n",
    "    llmModel = BedrockModels.AnthropicClaude35SonnetV2, // State-of-the-art reasoning model\n",
    "    systemPrompt = \"\"\"\n",
    "        You are a helpful assistant that controls a switch device.\n",
    "\n",
    "        You can:\n",
    "        - Turn the switch on or off when requested\n",
    "        - Check the current state of the switch\n",
    "        - Explain what you're doing\n",
    "\n",
    "        Always be clear about the switch's current state and confirm actions taken.\n",
    "    \"\"\".trimIndent(),\n",
    "    temperature = 0.1, // Low temperature for consistent, focused responses\n",
    "    toolRegistry = toolRegistry\n",
    ")\n",
    "\n",
    "println(\"🤖 AI Agent created successfully!\")\n",
    "println(\"📋 System prompt configured\")\n",
    "println(\"🛠️  Tools available: ${toolRegistry.tools.size}\")\n",
    "println(\"🎯 Model: ${BedrockModels.AnthropicClaude35SonnetV2}\")\n",
    "println(\"🌡️  Temperature: 0.1 (focused responses)\")"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖 AI Agent created successfully!\n",
      "📋 System prompt configured\n",
      "🛠️  Tools available: 2\n",
      "🎯 Model: LLModel(provider=Bedrock, id=us.anthropic.claude-3-5-sonnet-20241022-v2:0, capabilities=[Temperature, Tools, ToolChoice, Image, Document, Completion], contextLength=200000, maxOutputTokens=8192)\n",
      "🌡️  Temperature: 0.1 (focused responses)\n"
     ]
    }
   ],
   "execution_count": 6
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-08-18T21:36:26.995626Z",
     "start_time": "2025-08-18T21:36:23.860790Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import kotlinx.coroutines.runBlocking\n",
    "\n",
    "println(\"🎉 Bedrock Agent with Switch Tools - Ready to Go!\")\n",
    "println(\"💬 You can ask me to:\")\n",
    "println(\"   • Turn the switch on/off\")\n",
    "println(\"   • Check the current switch state\")\n",
    "println(\"   • Ask questions about the switch\")\n",
    "println()\n",
    "println(\"💡 Example: 'Please turn on the switch' or 'What's the current state?'\")\n",
    "println(\"📝 Type your request:\")\n",
    "\n",
    "val input = readln()\n",
    "println(\"\\n🤖 Processing your request...\")\n",
    "\n",
    "runBlocking {\n",
    "    val response = agent.run(input)\n",
    "    println(\"\\n✨ Agent response:\")\n",
    "    println(response)\n",
    "}"
   ],
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🎉 Bedrock Agent with Switch Tools - Ready to Go!\n",
      "💬 You can ask me to:\n",
      "   • Turn the switch on/off\n",
      "   • Check the current switch state\n",
      "   • Ask questions about the switch\n",
      "\n",
      "💡 Example: 'Please turn on the switch' or 'What's the current state?'\n",
      "📝 Type your request:\n"
     ]
    },
    {
     "ename": "org.jetbrains.kotlinx.jupyter.exceptions.ReplInterruptedException",
     "evalue": "The execution was interrupted",
     "output_type": "error",
     "traceback": [
      "The execution was interrupted"
     ]
    }
   ],
   "execution_count": 7
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## What Just Happened? 🎯\n",
    "\n",
    "When you run the agent, here's the magic that occurs behind the scenes:\n",
    "\n",
    "1. **Natural Language Processing**: Your input is sent to Claude 3.5 Sonnet via Bedrock\n",
    "2. **Intent Recognition**: The model understands what you want to do with the switch\n",
    "3. **Tool Selection**: Based on your request, the agent decides which tools to call\n",
    "4. **Action Execution**: The appropriate tool methods are invoked on your switch object\n",
    "5. **Response Generation**: The agent formulates a natural language response about what happened\n",
    "\n",
    "This demonstrates the core power of the Koog framework - seamless integration between natural language understanding and programmatic actions.\n",
    "\n",
    "## Next Steps & Extensions\n",
    "\n",
    "Ready to take this further? Here are some ideas to explore:\n",
    "\n",
    "### 🔧 Enhanced Tools\n",
    "```kotlin\n",
    "@Tool\n",
    "@LLMDescription(\"Sets a timer to automatically turn off the switch after specified seconds\")\n",
    "fun setAutoOffTimer(seconds: Int): String\n",
    "\n",
    "@Tool\n",
    "@LLMDescription(\"Gets the switch usage statistics and history\")\n",
    "fun getUsageStats(): String\n",
    "```\n",
    "\n",
    "### 🌐 Multiple Devices\n",
    "```kotlin\n",
    "class HomeAutomationTools : ToolSet {\n",
    "    @Tool fun controlLight(room: String, on: Boolean): String\n",
    "    @Tool fun setThermostat(temperature: Double): String\n",
    "    @Tool fun lockDoor(doorName: String): String\n",
    "}\n",
    "```\n",
    "\n",
    "### 🧠 Memory & Context\n",
    "```kotlin\n",
    "val agent = AIAgent(\n",
    "    executor = executor,\n",
    "    // ... other config\n",
    "    features = listOf(\n",
    "        MemoryFeature(), // Remember past interactions\n",
    "        LoggingFeature()  // Track all actions\n",
    "    )\n",
    ")\n",
    "```\n",
    "\n",
    "### 🔄 Advanced Workflows\n",
    "```kotlin\n",
    "// Multi-step workflows with conditional logic\n",
    "@Tool\n",
    "@LLMDescription(\"Executes evening routine: dims lights, locks doors, sets thermostat\")\n",
    "fun eveningRoutine(): String\n",
    "```\n",
    "\n",
    "## Key Takeaways\n",
    "\n",
    "✅ **Tools are functions**: Any Kotlin function can become an agent capability\n",
    "✅ **Annotations drive behavior**: @Tool and @LLMDescription make functions discoverable\n",
    "✅ **ToolSets organize capabilities**: Group related tools together logically\n",
    "✅ **Registries are toolboxes**: ToolRegistry contains all available agent capabilities\n",
    "✅ **Agents orchestrate everything**: AIAgent brings LLM intelligence + tools together\n",
    "\n",
    "The Koog framework makes it incredibly straightforward to build sophisticated AI agents that can understand natural language and take real-world actions. Start simple, then expand your agent's capabilities by adding more tools and features as needed.\n",
    "\n",
    "**Happy agent building!** 🚀"
   ]
  },
  {
   "metadata": {},
   "cell_type": "markdown",
   "source": [
    "## Testing the Agent\n",
    "\n",
    "Time to see our agent in action! The agent can now understand natural language requests and use the tools we've provided to control the switch.\n",
    "\n",
    "**Try these commands:**\n",
    "- \"Turn on the switch\"\n",
    "- \"What's the current state?\"\n",
    "- \"Switch it off please\"\n",
    "- \"Is the switch on or off?\""
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Kotlin",
   "language": "kotlin",
   "name": "kotlin"
  },
  "language_info": {
   "name": "kotlin",
   "version": "2.2.20-Beta2",
   "mimetype": "text/x-kotlin",
   "file_extension": ".kt",
   "pygments_lexer": "kotlin",
   "codemirror_mode": "text/x-kotlin",
   "nbconvert_exporter": ""
  },
  "ktnbPluginMetadata": {
   "projectDependencies": [
    "koog-agents.examples.main"
   ],
   "projectLibraries": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
